_base_ = ['./vidar_full_nusc_1future.py']

# nSKG配置
use_nskg = True
nskg_path = 'data/nuscenes/nskg/nSKG.ttl'
nskg_ontology_path = 'data/nuscenes/nskg/ontology'
use_nstp = False  # 是否使用nSTP格式

# 修改数据集配置
train_pipeline = _base_.train_pipeline
test_pipeline = _base_.test_pipeline

data = dict(
    samples_per_gpu=1,
    workers_per_gpu=4,
    train=dict(
        type='NuScenesViDARDatasetV1',
        use_nskg=use_nskg,
        nskg_path=nskg_path,
        nskg_ontology_path=nskg_ontology_path,
        use_nstp=use_nstp,
        # 其他参数保持不变
    ),
    val=dict(
        type='NuScenesViDARDatasetV1',
        use_nskg=use_nskg,
        nskg_path=nskg_path,
        nskg_ontology_path=nskg_ontology_path,
        use_nstp=use_nstp,
        # 其他参数保持不变
    ),
    test=dict(
        type='NuScenesViDARDatasetV1',
        use_nskg=use_nskg,
        nskg_path=nskg_path,
        nskg_ontology_path=nskg_ontology_path,
        use_nstp=use_nstp,
        # 其他参数保持不变
    ),
)

# 修改模型配置
model = dict(
    pts_bbox_head=dict(
        type='ViDARHeadV1',
        use_nskg=use_nskg,
        nskg_encoder_cfg=dict(
            in_channels=8,
            hidden_channels=64,
            out_channels=256,
            num_layers=2,
            gnn_type='gat',
            use_hetero=True,
            dropout=0.1
        ),
        nskg_enhancer_cfg=dict(
            bev_channels=256,  # 与embed_dims保持一致
            nskg_channels=256,
            hidden_channels=128,
            bev_h=200,  # 与bev_h保持一致
            bev_w=200,  # 与bev_w保持一致
            use_attention=True
        ),
        # 其他参数保持不变
    )
)

# 优化器配置 - 可能需要调整学习率
optimizer = dict(
    type='AdamW',
    lr=2e-4,  # 可能需要降低学习率以适应新增的nSKG模块
    weight_decay=0.01,
)